# import os
# import json
#
# import torch
# from PIL import Image
# from torchvision import transforms
# import matplotlib.pyplot as plt
#
# from vit_model import ViTImageClassifier as create_model
#
#
# def main():
#     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#
#     data_transform = transforms.Compose(
#         [transforms.Resize(256),
#          transforms.CenterCrop(224),
#          transforms.ToTensor(),
#          transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
#
#     # load image
#     img_path = r"E:\ViT\datasets\cat_dog\dog\dog.2.jpg"
#     assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
#     img = Image.open(img_path)
#     plt.imshow(img)
#     # [N, C, H, W]
#     img = data_transform(img)
#     # expand batch dimension
#     img = torch.unsqueeze(img, dim=0)
#
#     # read class_indict
#     json_path = 'E:\ViT\class_indices.json'
#     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
#
#     with open(json_path, "r") as f:
#         class_indict = json.load(f)
#
#     # create model
#     model = create_model(num_classes=2, has_logits=False).to(device)
#     # load model weights
#     model_weight_path = r"E:\ViT\cat_dog.pth"
#     # model_weight_path = r"E:\ViT\output\super_baby_checkpoint.bin"
#     model.load_state_dict(torch.load(model_weight_path, map_location=device))
#     model.eval()
#     with torch.no_grad():
#         # predict class
#         output = torch.squeeze(model(img.to(device))).cpu()
#         predict = torch.softmax(output, dim=0)
#         predict_cla = torch.argmax(predict).numpy()
#
#     print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
#                                                  predict[predict_cla].numpy())
#     plt.title(print_res)
#     for i in range(len(predict)):
#         print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
#                                                   predict[i].numpy()))
#     plt.show()
#
#
# if __name__ == '__main__':
#     main()

import os
import json
import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
from transformers import ViTForImageClassification, ViTConfig


class ViTImageClassifier:
    def __init__(self, config):
        # 初始化设备
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        # 图像预处理
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

        # 初始化模型
        self.model = ViTForImageClassification(config).to(self.device)
        self.class_indices = None

    def load_model_weights(self, weight_path):
        """加载模型权重"""
        try:
            state_dict = torch.load(weight_path, map_location=self.device)
            self.model.load_state_dict(state_dict)
            self.model.eval()
            return True
        except Exception as e:
            print(f"加载模型权重失败: {e}")
            return False

    def load_class_indices(self, json_path):
        """加载类别索引"""
        try:
            with open(json_path, "r") as f:
                self.class_indices = json.load(f)
            return True
        except Exception as e:
            print(f"加载类别索引失败: {e}")
            return False

    def predict(self, img_path, show_result=True):
        """执行预测"""
        # 验证文件存在
        if not os.path.exists(img_path):
            print(f"文件不存在: {img_path}")
            return None

        # 加载并预处理图像
        try:
            img = Image.open(img_path)
            img_tensor = self.transform(img).unsqueeze(0).to(self.device)
        except Exception as e:
            print(f"图像加载/预处理失败: {e}")
            return None

        # 执行预测
        with torch.no_grad():
            try:
                outputs = self.model(img_tensor)
                probs = torch.softmax(outputs.logits.squeeze(), dim=0)
                pred_class = torch.argmax(probs).item()

                # 显示结果
                if show_result:
                    self._display_result(img, pred_class, probs)

                return {
                    "class": self.class_indices[str(pred_class)],
                    "probability": probs[pred_class].item(),
                    "all_probs": {self.class_indices[str(i)]: probs[i].item()
                                  for i in range(len(probs))}
                }
            except Exception as e:
                print(f"预测失败: {e}")
                return None

    def _display_result(self, img, pred_class, probs):
        """可视化显示结果"""
        plt.imshow(img)
        title = f"预测类别: {self.class_indices[str(pred_class)]}  概率: {probs[pred_class]:.3f}"
        plt.title(title)

        # 打印所有类别概率
        print("\n预测结果:")
        for i, prob in enumerate(probs):
            print(f"{self.class_indices[str(i)]:>15}: {prob:.3f}")

        plt.show()


def main():
    # 配置模型参数
    config = ViTConfig()
    config.num_classes = 2  # 根据你的分类任务调整

    # 初始化分类器
    classifier = ViTImageClassifier(config)

    # 加载模型权重
    model_path = r"E:\ViT\cat_dog.pth"
    if not classifier.load_model_weights(model_path):
        return

    # 加载类别索引
    class_json = r"E:\ViT\class_indices.json"
    if not classifier.load_class_indices(class_json):
        return

    # 执行预测
    img_path = r"E:\ViT\datasets\cat_dog\dog\dog.2.jpg"
    result = classifier.predict(img_path)

    if result:
        print("\n最终预测结果:")
        print(f"类别: {result['class']}")
        print(f"置信度: {result['probability']:.2%}")


if __name__ == '__main__':
    main()








